package com.xrui.hbase.impl;

import com.google.common.base.Preconditions;

import com.xrui.hbase.*;
import com.xrui.hbase.codec.impl.CellDecoderProvider;
import com.xrui.hbase.util.DebugResourceTracker;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.regionserver.LeaseException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.time.Instant;
import java.util.Arrays;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.concurrent.atomic.AtomicReference;


/**
 * The internal implementation of RowScanner that reads from HTables.
 */
public final class HBaseRowScanner implements RowScanner {
    public static final String MAX_RETRIES_ON_TIMEOUT_PROPERTY =
        "com.tr.ts.hbase.impl.hbase.HBaseRowScanner.MAX_RETRIES_ON_TIMEOUT";
    private static final Logger LOG = LoggerFactory.getLogger(HBaseRowScanner.class);
    /**
     * Maximum number of retries when (re-)opening the HBase scanner.
     */
    private static final int MAX_RETRIES_ON_TIMEOUT =
        Integer.parseInt(System.getProperty(MAX_RETRIES_ON_TIMEOUT_PROPERTY, "3"));

    /**
     * The request used to fetch the row data.
     */
    private final DataRequest mDataRequest;

    /**
     * The table being scanned.
     */
    private final HBaseTable mTable;

    /**
     * HBase scan specification.
     */
    private final Scan mScan;

    /**
     * Provider for cell decoders.
     */
    private final CellDecoderProvider mCellDecoderProvider;
    /**
     * Tracks the state of this row scanner.
     */
    private final AtomicReference<State> mState = new AtomicReference<State>(State.UNINITIALIZED);
    /**
     * Whether to reopen the HBase scanner on timeout.
     */
    private final boolean mReopenScannerOnTimeout;
    /**
     * HTable connection.
     */
    private final HTableInterface mHTable;
    /**
     * Current HBase result scanner. This scanner may timeout.
     */
    private ResultScanner mResultScanner = null;
    /**
     * Result to return to the user on the following invocation of next().
     */
    private Result mNextResult = null;
    /**
     * HBase row key of the last result returned to the user.
     */
    private byte[] mLastReturnedKey = null;

    /**
     * Creates a new <code>HBaseRowScanner</code> instance.
     *
     * @param options The options for this scanner.
     * @throws IOException on I/O error.
     */
    public HBaseRowScanner(Options options) throws IOException {
        mDataRequest = options.getDataRequest();
        mTable = options.getTable();
        mScan = options.getScan();
        mCellDecoderProvider = options.getCellDecoderProvider();
        mReopenScannerOnTimeout = options.getReopenScannerOnTimeout();

        mHTable = mTable.openHTableConnection();
        try {
            mResultScanner = openResultScanner();
            mNextResult = getNextResult();
        } catch (IOException ioe) {
            if (mHTable != null) {
                mHTable.close();
            }
            throw ioe;
        }

        final State oldState = mState.getAndSet(State.OPEN);
        Preconditions.checkState(oldState == State.UNINITIALIZED,
            "Cannot open RowScanner instance in state %s.", oldState);
        DebugResourceTracker.get().registerResource(this);
    }

    // -----------------------------------------------------------------------------------------------

    /**
     * Computes the smallest possible HBase row key strictly greater than a given row key.
     *
     * @param rowKey A given HBase row key.
     * @return the smallest possible HBase row key strictly greater than the given row key.
     */
    private static byte[] getSmallestHigherThan(byte[] rowKey) {
        return Arrays.copyOf(rowKey, rowKey.length + 1);
    }

    // -----------------------------------------------------------------------------------------------

    /**
     * Opens a new HBase scanner.
     * <p>
     * <p> Adjusts the start scanning row based on the last row returned to the user. </p>
     *
     * @return a new HBase scanner.
     */
    private ResultScanner openResultScanner() throws IOException {
        try {
            if (mLastReturnedKey != null) {
                // If we previously returned a row to the user,
                // start the new scan at the lowest possible next row:
                mScan.setStartRow(getSmallestHigherThan(mLastReturnedKey));
            }
            LOG.debug("Opening HBase result scanner with scan option: '{}'.", mScan.toString());

            return mHTable.getScanner(mScan);
        } catch (IOException ioe) {
            throw new IOException(ioe);
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public HBaseRowIterator iterator() {
        return new HBaseRowIterator();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void close() throws IOException {
        final State oldState = mState.getAndSet(State.CLOSED);
        Preconditions.checkState(oldState == State.OPEN,
            "Cannot close HBaseRowScanner instance in state %s.", oldState);
        DebugResourceTracker.get().unregisterResource(this);
        mResultScanner.close();
        mHTable.close();
    }

    /**
     * Fetches the next HBase result.
     * <p>
     * <p> Handles HBase scanner timeouts. </p>
     *
     * @return the next HBase result, or null if none.
     */
    private Result getNextResult() throws IOException {
        for (int nretries = 0; nretries < MAX_RETRIES_ON_TIMEOUT; ++nretries) {
            try {
                return mResultScanner.next();

            } catch (LeaseException le) {
                if (!mReopenScannerOnTimeout) {
                    LOG.debug("HBase scanner timed out and user disabled automatic scanner reopening.");
                    throw new IOException(
                        "HBase scanner timed out and user disabled automatic scanner reopening.", le);
                } else {
                    // The HBase scanner timed out, re-open a new one:
                    LOG.debug("HBase scanner timed out: closing and reopening a new scanner.");
                    mResultScanner.close();
                    LOG.debug("scan begin" + Instant.now().toString());
                    mResultScanner = openResultScanner();
                    LOG.debug("scan end" + Instant.now().toString());
                    continue;
                }

            } catch (ScannerTimeoutException ste) {
                if (!mReopenScannerOnTimeout) {
                    LOG.debug("HBase scanner timed out and user disabled automatic scanner reopening.");
                    throw new IOException(
                        "HBase scanner timed out and user disabled automatic scanner reopening.", ste);
                } else {
                    // The HBase scanner timed out, re-open a new one:
                    LOG.debug("HBase scanner timed out: closing and reopening a new scanner.");
                    mResultScanner.close();
                    mResultScanner = openResultScanner();
                    continue;
                }

            } catch (IOException ioe) {
                throw new IOException(ioe);
            }
        }
        throw new IOException("Unable to retrieve HBase result from scanner.");
    }

    /**
     * States of a row scanner instance.
     */
    private static enum State {
        UNINITIALIZED,
        OPEN,
        CLOSED
    }

    /**
     * A class to encapsulate the various options the HBaseRowScanner constructor requires.
     */
    public static class Options {
        private DataRequest mDataRequest;
        private HBaseTable mTable;
        private Scan mScan;
        private CellDecoderProvider mCellDecoderProvider;
        private boolean mReopenScannerOnTimeout;

        /**
         * Sets the data request used to generate the RowScanner.
         *
         * @param dataRequest A data request.
         * @return This options instance.
         */
        public Options withDataRequest(DataRequest dataRequest) {
            mDataRequest = dataRequest;
            return this;
        }

        /**
         * Sets the table being scanned.
         *
         * @param table The table being scanned.
         * @return This options instance.
         */
        public Options withTable(HBaseTable table) {
            mTable = table;
            return this;
        }

        /**
         * Sets the HBase scan specification.
         *
         * @param scan HBase scan specification.
         * @return This options instance.
         */
        public Options withScan(Scan scan) {
            mScan = scan;
            return this;
        }

        /**
         * Sets whether the HBase scanner should be reopened on timeout.
         *
         * @param reopenScannerOnTimeout Whether to reopen the HBase scanner on timeout.
         * @return This options instance.
         */
        public Options withReopenScannerOnTimeout(boolean reopenScannerOnTimeout) {
            mReopenScannerOnTimeout = reopenScannerOnTimeout;
            return this;
        }

        /**
         * Sets a provider for cell decoders.
         *
         * @param cellDecoderProvider Provider for cell decoders.
         * @return This options instance.
         */
        public Options withCellDecoderProvider(CellDecoderProvider cellDecoderProvider) {
            mCellDecoderProvider = cellDecoderProvider;
            return this;
        }

        /**
         * Gets the data request.
         *
         * @return The data request.
         */
        public DataRequest getDataRequest() {
            return mDataRequest;
        }

        /**
         * Gets the table being scanned.
         *
         * @return The table.
         */
        public HBaseTable getTable() {
            return mTable;
        }

        /**
         * Gets the HBase scan specification.
         *
         * @return the HBase scan specification.
         */
        public Scan getScan() {
            return mScan;
        }

        /**
         * Gets the provider for cell decoders.
         *
         * @return the provider for cell decoders.
         */
        public CellDecoderProvider getCellDecoderProvider() {
            return mCellDecoderProvider;
        }

        /**
         * Reports whether the HBase scanner should be re-opened on timeout.
         *
         * @return whether the HBase scanner should be re-opened on timeout.
         */
        public boolean getReopenScannerOnTimeout() {
            return mReopenScannerOnTimeout;
        }
    }

    // -----------------------------------------------------------------------------------------------

    /**
     * Wraps a row scanner into a Java iterator.
     */
    private class HBaseRowIterator implements Iterator<RowData> {
        /**
         * {@inheritDoc}
         */
        @Override
        public boolean hasNext() {
            final State state = mState.get();
            Preconditions.checkState(state == State.OPEN,
                "Cannot check has next on RowScanner instance in state %s.", state);
            return (mNextResult != null);
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public RowData next() {
            final State state = mState.get();
            Preconditions.checkState(state == State.OPEN,
                "Cannot create next on RowScanner instance in state %s.", state);
            if (mNextResult == null) {
                // Comply with the Iterator interface:
                throw new NoSuchElementException();
            }
            final Result result = mNextResult;
            mLastReturnedKey = result.getRow();

            try {
                // Prefetch the next row for hasNext():
                mNextResult = getNextResult();
                // Decode the HBase result into a RowData:

                final EntityId entityId = HBaseEntityId.create(result.getRow(), mTable.getLayout());
                return new HBaseRowData(mTable, mDataRequest, entityId, result, mCellDecoderProvider);
            } catch (IOException ioe) {
                throw new com.xrui.hbase.exception.IOException(ioe);
            }
        }

        /**
         * {@inheritDoc}
         */
        @Override
        public void remove() {
            throw new UnsupportedOperationException("RowIterator does not support remove().");
        }
    }
}
